﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Reflection;
using System.IO;
using Microsoft.Win32;
using Kokomo.Mirror.PeFormat;

namespace Kokomo.Mirror.CliMetadata
{
	public class MirrorLoadContext : IAssemblyLoadContext
	{
		public static MirrorLoadContext DefaultLoadContext { get; private set; }
		static MirrorLoadContext()
		{
			DefaultLoadContext = new MirrorLoadContext(null);

			Dictionary<TypeCode, string> typeCodeNameTable = new Dictionary<TypeCode, string>();
			typeCodeNameTable[TypeCode.Object] = typeof(object).FullName;
			typeCodeNameTable[TypeCode.DBNull] = typeof(DBNull).FullName;
			typeCodeNameTable[TypeCode.Boolean] = typeof(bool).FullName;
			typeCodeNameTable[TypeCode.Char] = typeof(char).FullName;
			typeCodeNameTable[TypeCode.SByte] = typeof(sbyte).FullName;
			typeCodeNameTable[TypeCode.Byte] = typeof(byte).FullName;
			typeCodeNameTable[TypeCode.Int16] = typeof(short).FullName;
			typeCodeNameTable[TypeCode.UInt16] = typeof(ushort).FullName;
			typeCodeNameTable[TypeCode.Int32] = typeof(int).FullName;
			typeCodeNameTable[TypeCode.UInt32] = typeof(uint).FullName;
			typeCodeNameTable[TypeCode.Int64] = typeof(long).FullName;
			typeCodeNameTable[TypeCode.UInt64] = typeof(ulong).FullName;
			typeCodeNameTable[TypeCode.Single] = typeof(float).FullName;
			typeCodeNameTable[TypeCode.Double] = typeof(double).FullName;
			typeCodeNameTable[TypeCode.Decimal] = typeof(decimal).FullName;
			typeCodeNameTable[TypeCode.DateTime] = typeof(DateTime).FullName;
			typeCodeNameTable[TypeCode.String] = typeof(string).FullName;

			_typeCodeNameTable = typeCodeNameTable;
			_typeCodeIndexTable = typeCodeNameTable.ToDictionary(e => e.Value, e => e.Key);
		}

		private static Dictionary<string, TypeCode> _typeCodeIndexTable;
		private static Dictionary<TypeCode, string> _typeCodeNameTable;

		private List<Assembly> _loadedAssemblies;
		public IEnumerable<Assembly> LoadedAssemblies { get { return this._loadedAssemblies; } }

		private Dictionary<string, CliAssembly> _assemblyFileTable;
		//private Dictionary<AssemblyName, CliAssembly> _assemblyNameTable;

		public string PrimarySearchPath { get; private set; }

		public MirrorLoadContext(
			string primarySearchPath
			)
		{
			//if (primarySearchPath == null) throw new ArgumentNullException("primarySearchPath");
			this.PrimarySearchPath = primarySearchPath;
			this._loadedAssemblies = new List<Assembly>();
			this._assemblyFileTable = new Dictionary<string, CliAssembly>(StringComparer.InvariantCultureIgnoreCase);
			//this._assemblyNameTable = new Dictionary<AssemblyName, CliAssembly>(new AssemblyNameComparer());

			this._typeCodeTable = new Dictionary<TypeCode, Type>();

			this._legacyGacRootPath = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.Windows), "assembly");
			this._gacRootPath = Path.Combine(Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.Windows), "Microsoft.NET"), "assembly");
			this._commonLanguageRuntimeLibrary = new Lazy<Assembly>(this.LoadCommonLanguageRuntimeLibrary);

			this._valueType = new Lazy<Type>(this.LoadValueType);
			this._intPtrType = new Lazy<Type>(this.LoadIntPtrType);
		}

		private string _legacyGacRootPath;
		private string _gacRootPath;

		private Lazy<Assembly> _commonLanguageRuntimeLibrary;
		private Assembly LoadCommonLanguageRuntimeLibrary()
		{
			return this.LoadAssemblyFromGac(new AssemblyName() { Name = "mscorlib" });
		}
		public Assembly CommonLanguageRuntimeLibrary { get { return this._commonLanguageRuntimeLibrary.Value; } }

		#region System types
		private Dictionary<TypeCode, Type> _typeCodeTable;
		public Type GetType(TypeCode typeCode)
		{
			lock (this)
			{
				Type type;
				if (this._typeCodeTable.TryGetValue(typeCode, out type))
					return type;

				string typeName;
				if (_typeCodeNameTable.TryGetValue(typeCode, out typeName))
				{
					type = this.CommonLanguageRuntimeLibrary.GetType(typeName);
					this._typeCodeTable.Add(typeCode, type);
				}
				else
				{
					throw new ArgumentException("Invalid TypeCode", "TypeCode");
				}

				return type;
			}
		}
		internal static TypeCode GetTypeCode(Type type)
		{
			TypeCode code;
			if (!_typeCodeIndexTable.TryGetValue(type.FullName, out code))
				code = TypeCode.Object;

			return code;
		}

		public Type GetType(string typeName)
		{
			return Type.GetType(typeName, this.LoadAssembly, null);
		}

		private Lazy<Type> _valueType;
		private Type LoadValueType()
		{
			return this.CommonLanguageRuntimeLibrary.GetType("System.ValueType");
		}
		public Type ValueType { get { return this._valueType.Value; } }

		private Lazy<Type> _intPtrType;
		private Type LoadIntPtrType()
		{
			return this.CommonLanguageRuntimeLibrary.GetType("System.IntPtr");
		}
		public Type IntPtrType { get { return this._intPtrType.Value; } }
		#endregion

		public CliAssembly GetLoadedAssembly(AssemblyName assemblyName)
		{
			if (assemblyName == null) throw new ArgumentNullException("assemblyName");

			// The extra ToArray() protects against the collection being modified during enumeration
			CliAssembly assembly = (CliAssembly)this._loadedAssemblies.ToArray().FirstOrDefault(
				a => AssemblyName.ReferenceMatchesDefinition(assemblyName, a.GetName())
					);
			//this._assemblyNameTable.TryGetValue(assemblyName, out assembly);
			return assembly;
		}
		private void OnAssemblyLoaded(CliAssembly assembly)
		{
			if (assembly == null) throw new ArgumentNullException("assembly");

			AssemblyName assemblyName = assembly.GetName();
			this._loadedAssemblies.Add(assembly);
			this._assemblyFileTable[assembly.Location] = assembly;
			//this._assemblyNameTable[assembly.GetName()] = assembly;
		}

		public CliAssembly LoadAssembly(AssemblyName assemblyName)
		{
			if (assemblyName == null) throw new ArgumentNullException("assemblyName");

			CliAssembly assembly = this.GetLoadedAssembly(assemblyName);
			if ((assembly == null) && (this.PrimarySearchPath != null))
				assembly = this.LoadAssembly(assemblyName, this.PrimarySearchPath);
			if (assembly == null)
				assembly = this.LoadAssemblyFromGac(assemblyName);
			if (assembly == null)
				assembly = this.LoadAssemblyFromLegacyGac(assemblyName);
			if ((assembly == null) && assemblyName.Version != null)
				assembly = this.LoadAssembly(new AssemblyName(assemblyName.Name));

			return assembly;
		}

		private CliAssembly LoadAssemblyFromGac(
			string gacRootPath,
			string principalName,
			string versionLabel
			)
		{
			if (gacRootPath == null) throw new ArgumentNullException("gacRootPath");
			if (principalName == null) throw new ArgumentNullException("principalName");

			string[] gacCategories = new string[] { "GAC", "GAC_32", "GAC_MSIL" };

			CliAssembly assembly = null;
			foreach (var gacCategory in gacCategories)
			{
				if (versionLabel != null)
				{
					string fullPath = Path.Combine(gacRootPath, gacCategory, principalName, versionLabel, principalName + ".DLL");
					if (File.Exists(fullPath))
					{
						assembly = this.LoadAssemblyFile(fullPath);
						if (assembly != null) break;
					}
				}
				else
				{
					string directoryPath = Path.Combine(gacRootPath, gacCategory, principalName);
					if (Directory.Exists(directoryPath))
					{
						string versionDirectory = Directory.GetDirectories(directoryPath).FirstOrDefault();
						if (versionDirectory != null)
						{
							string fullPath = Path.Combine(versionDirectory, principalName + ".DLL");
							if (File.Exists(fullPath))
							{
								assembly = this.LoadAssemblyFile(fullPath);
								if (assembly != null) break;
							}
						}
					}
				}
			}

			return assembly;
		}

		private CliAssembly LoadAssemblyFromLegacyGac(AssemblyName assemblyName)
		{
			if (assemblyName == null) throw new ArgumentNullException("assemblyName");

			string principalName = assemblyName.Name;
			string versionLabel = null;

			byte[] publicKeyToken = assemblyName.GetPublicKeyToken();
			if (publicKeyToken != null)
			{
				StringBuilder publicKeyTokenBuilder = new StringBuilder(16);
				foreach (var b in publicKeyToken)
				{
					publicKeyTokenBuilder.Append(b.ToString("x2"));
				}
				versionLabel = assemblyName.Version.ToString() + "__" + publicKeyTokenBuilder.ToString();
			}

			CliAssembly assembly = this.LoadAssemblyFromGac(
				this._legacyGacRootPath,
				principalName,
				versionLabel
				);
			return assembly;
		}

		private CliAssembly LoadAssemblyFromGac(AssemblyName assemblyName)
		{
			if (assemblyName == null) throw new ArgumentNullException("assemblyName");

			byte[] publicKeyToken = assemblyName.GetPublicKeyToken();

			string principalName = assemblyName.Name;
			string versionLabel = null;
			if (publicKeyToken != null)
			{
				StringBuilder publicKeyTokenBuilder = new StringBuilder(16);
				foreach (var b in publicKeyToken)
				{
					publicKeyTokenBuilder.Append(b.ToString("x2"));
				}
				versionLabel = "v4.0_" + assemblyName.Version.ToString() + "__" + publicKeyTokenBuilder.ToString();
			}

			CliAssembly assembly = this.LoadAssemblyFromGac(
				this._gacRootPath,
				principalName,
				versionLabel
				);
			return assembly;
		}
		public CliAssembly LoadAssembly(AssemblyName assemblyName, string searchPath)
		{
			if (assemblyName == null) throw new ArgumentNullException("assemblyName");
			if (searchPath == null) throw new ArgumentNullException("searchPath");

			// Determine file name
			string assemblyFileName = assemblyName.Name + ".DLL";
			string assemblyPath = Path.Combine(searchPath, assemblyFileName);

			if (!File.Exists(assemblyPath)) return null;

			CliAssembly assembly = this.LoadAssemblyFile(assemblyPath);

			return assembly;
		}

		public CliAssembly LoadAssemblyFile(string assemblyPath)
		{
			if (assemblyPath == null) throw new ArgumentNullException("assemblyPath");

			CliAssembly assembly;
			lock (this)
			{
				if (this._assemblyFileTable.TryGetValue(assemblyPath, out assembly)) return assembly;

				PeModule pe = PeModule.Open(assemblyPath, PeModuleLoadOptions.LoadModule);
				assembly = new CliAssembly(this, pe);

				this.OnAssemblyLoaded(assembly);
			}

			return assembly;
		}
		public CliAssembly LoadAssembly(AssemblyName assemblyName, Assembly referringAssembly)
		{
			if (assemblyName == null) throw new ArgumentNullException("assemblyName");
			if (referringAssembly == null) throw new ArgumentNullException("referringAssembly");

			CliAssembly assembly = this.GetLoadedAssembly(assemblyName);
			if (assembly == null)
			{
				string localPath = Path.GetDirectoryName(referringAssembly.Location);
				assembly = this.LoadAssembly(assemblyName, localPath);
			}
			//if (assembly == null) assembly = this.LoadAssemblyFromGac(assemblyName);
			if (assembly == null) assembly = this.LoadAssembly(assemblyName);

			return assembly;
		}

		Assembly IAssemblyLoadContext.LoadAssembly(AssemblyName assemblyName)
		{
			return this.LoadAssembly(assemblyName);
		}

		Assembly IAssemblyLoadContext.LoadAssembly(AssemblyName assemblyName, Assembly referencingAssembly)
		{
			return this.LoadAssembly(assemblyName, referencingAssembly);
		}

		Assembly IAssemblyLoadContext.LoadAssemblyFile(string assemblyPath)
		{
			return this.LoadAssemblyFile(assemblyPath);
		}

		Assembly IAssemblyLoadContext.LoadAssembly(AssemblyName assemblyName, string searchDirectory)
		{
			return this.LoadAssembly(assemblyName, searchDirectory);
		}

		private class AssemblyNameComparer : IEqualityComparer<AssemblyName>
		{
			public bool Equals(AssemblyName x, AssemblyName y)
			{
				bool equals =
					(x.Name == y.Name)
					&& (
						(x.Version == y.Version)
						|| (x.Version == null)
						|| (y.Version == null)
						)
					;

				return equals;
			}

			public int GetHashCode(AssemblyName obj)
			{
				if (obj == null) return 0;

				int hash = obj.Name.GetHashCode();
				if (obj.Version != null) hash ^= obj.Version.GetHashCode();

				return hash;
			}
		}

		public Assembly[] GetLoadedAssemblies()
		{
			return this.LoadedAssemblies.ToArray();
		}
	}
}
